import pefile
from capstone import Cs, CS_ARCH_X86, CS_MODE_32

class StaticAnalyzer:
    def __init__(self, file_path):
        self.pe = pefile.PE(file_path)
        self.md = Cs(CS_ARCH_X86, CS_MODE_32)
        
    def extract_pe_features(self):
        """提取增强的PE特征"""
        features = {
            # PE头信息
            'header': {
                'machine': self.pe.FILE_HEADER.Machine,
                'timestamp': self.pe.FILE_HEADER.TimeDateStamp,
                'sections': len(self.pe.sections),
            },
            # 节表信息
            'sections': [{
                'name': section.Name.decode().rstrip('\x00'),
                'virtual_size': section.Misc_VirtualSize,
                'entropy': section.get_entropy(),
                'characteristics': section.Characteristics
            } for section in self.pe.sections],
            # 导入表分析
            'imports': {
                'dll_count': len(self.pe.DIRECTORY_ENTRY_IMPORT),
                'api_count': sum(len(imp.imports) for imp in self.pe.DIRECTORY_ENTRY_IMPORT)
            },
            # 操作码序列
            'opcodes': self._extract_opcodes()
        }
        return features
    
    def _extract_opcodes(self, max_length=500):
        """提取操作码序列"""
        opcodes = []
        for section in self.pe.sections:
            if section.Characteristics & 0x20000000:  # 可执行段
                data = section.get_data()
                for insn in self.md.disasm(data, 0x1000):
                    opcodes.append(insn.mnemonic)
                    if len(opcodes) >= max_length:
                        return opcodes[:max_length]
        return opcodes 